import fitz
import os
import sys
import PyPDF2
from Crypto.Cipher import AES
from multiprocessing import Pool
import psutil

class RmOne():
    def __init__(self,old) -> None:
        self.old = old.split(".pdf")[0]
        self.new = self.old + "new"
        self.new_file = self.new + ".pdf"
        self.old_file = old
    
    def run(self):
        pdf_reader = PyPDF2.PdfReader(self.old_file)
        total_pages = pdf_reader.getNumPages()
        pdf_writer = PyPDF2.PdfWriter()

        delete_pages = [0,]
        for i in range(total_pages):
            if i not in delete_pages:
                page = pdf_reader.pages[i]
                pdf_writer.add_page(page)

        with open(self.new_file,'wb') as f:
            pdf_writer.write(f)

        try:
            doc = fitz.open(self.new_file)
            doc.set_toc(self.get_toc())
            doc.saveIncr()
            doc.close()
        except ValueError:
            pass
        # os.system("mv '{}' '{}'".format(self.new_file,self.old_file))
        print("{} 处理完成".format(self.old_file))
    
    def get_toc(self):
        doc = fitz.open(self.old_file)
        toc = doc.get_toc()
        for i in toc :
            i[2] -=1
        doc.close
        return toc

def run(c):
    RmOne(c).run()

if __name__ == "__main__":
    file_names = os.listdir("./")
    file_names_new = []
    for i in file_names:
        if i.endswith(".pdf") or i.endswith(".PDF"):
            file_names_new.append(i)
    num_cpu = psutil.cpu_count()
    p = Pool(num_cpu)
    p.map(run,file_names_new)
